use std::{
    sync::Arc,
    net::ToSocketAddrs as _,
};
use anyhow::Error;
use quinn::*;
use tracing::*;
use tracing_subscriber::prelude::*;


#[tokio::main]
async fn main() -> Result<(), Error> {
    // init logging
    let log_fmt = tracing_subscriber::fmt::format()
        .compact()
        .with_timer(tracing_subscriber::fmt::time::uptime())
        .with_line_number(true);
    let stdout_log = tracing_subscriber::fmt::layer()
        .event_format(log_fmt)
        .with_writer(std::io::stderr);
    let log_filter = tracing_subscriber::EnvFilter::new(
        std::env::var(tracing_subscriber::EnvFilter::DEFAULT_ENV).unwrap_or("info".into())
    );
    let log_subscriber = tracing_subscriber::Registry::default()
        .with(log_filter)
        .with(stdout_log);
    tracing::subscriber::set_global_default(log_subscriber).expect("unable to install logger");

    // get args
    let args = std::env::args().collect::<Vec<_>>();
    anyhow::ensure!(args.len() == 2, "wrong number of args");
    let num_bytes = args[1].parse::<u32>()?;

    // generate keys
    let rcgen_cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
    let key = rustls::pki_types::PrivatePkcs8KeyDer::from(rcgen_cert.key_pair.serialize_der());
    let cert = rustls::pki_types::CertificateDer::from(rcgen_cert.cert);
    let mut roots = rustls::RootCertStore::empty();
    roots.add(cert.clone()).unwrap();
    let certs = vec![cert];

    let mut tasks = tokio::task::JoinSet::new();

    // start server
    let (send_stop_server, mut recv_stop_server) = tokio::sync::oneshot::channel();
    tasks.spawn(log_err(async move {
        let mut server_crypto = rustls::ServerConfig::builder()
                .with_no_client_auth()
                .with_single_cert(certs, key.into())?;
        // make sure to configure this:
        server_crypto.max_early_data_size = u32::MAX;
        let server_crypto = quinn::crypto::rustls::QuicServerConfig::try_from(Arc::new(server_crypto))?;
        let server_config = ServerConfig::with_crypto(Arc::new(server_crypto));
        let endpoint = Endpoint::server(
            server_config,
            "127.0.0.1:4433".to_socket_addrs().unwrap().next().unwrap(),
        )?;
        loop {
            let incoming = tokio::select! {
                option = endpoint.accept() => match option { Some(incoming) => incoming, None => break },
                result = &mut recv_stop_server => if result.is_ok() { break } else { continue },
            };
            // spawn subtask for connection
            tokio::spawn(log_err(async move {
                // attempt to accept 0-RTT data
                let conn = match incoming.accept()?.into_0rtt() {
                    Ok((conn, _)) => conn,
                    Err(connecting) => connecting.await?,
                };
                loop {
                    let (mut send, mut recv) = match conn.accept_bi().await {
                        Ok(stream) => stream,
                        Err(ConnectionError::ApplicationClosed(_)) => break,
                        Err(e) => Err(e)?,
                    };
                    // spawn subtask for stream
                    tokio::spawn(log_err(async move {
                        let requested_len_le_vec = recv.read_to_end(4).await?;
                        anyhow::ensure!(requested_len_le_vec.len() == 4, "malformed request {:?}", requested_len_le_vec);
                        let mut requested_len_le = [0; 4];
                        requested_len_le.copy_from_slice(&requested_len_le_vec);
                        let requested_len = u32::from_le_bytes(requested_len_le) as usize;
                        info!(%requested_len, "received request");
                        const BUF_LEN: usize = 8 << 10;
                        let mut buf = [0; BUF_LEN];
                        for i in 0..requested_len {
                            buf[i % BUF_LEN] = (i % 0xff) as u8;
                            if i % BUF_LEN == BUF_LEN - 1 {
                                send.write_all(&buf).await?;
                            }
                        }
                        if requested_len % BUF_LEN != 0 {
                            send.write_all(&buf[..requested_len % BUF_LEN]).await?;
                        }
                        info!("wrote response");
                        Ok(())
                    }.instrument(info_span!("server stream"))));
                }
                Ok(())
            }.instrument(info_span!("server conn"))));
        }
        // shut down server endpoint cleanly
        endpoint.wait_idle().await;
        Ok(())
    }.instrument(info_span!("server"))));

    // start client
    async fn send_request(conn: &Connection, num_bytes: u32) -> Result<std::time::Duration, Error> {
        let (mut send, mut recv) = conn.open_bi().await?;

        let start_time = std::time::Instant::now();

        debug!("sending request");
        send.write_all(&num_bytes.to_le_bytes()).await?;
        send.finish()?;
        debug!("receiving response");
        let response = recv.read_to_end(num_bytes as _).await?;
        anyhow::ensure!(response.len() == num_bytes as usize, "response is the wrong number of bytes");
        debug!("response received");

        let end_time = std::time::Instant::now();
        Ok(end_time.duration_since(start_time))
    }
    tasks.spawn(log_err(async move {
        let mut client_crypto = rustls::ClientConfig::builder()
                .with_root_certificates(roots)
                .with_no_client_auth();
        // make sure to configure this:
        client_crypto.enable_early_data = true;
        let mut endpoint = Endpoint::client(
            "0.0.0.0:0".to_socket_addrs().unwrap().next().unwrap()
        )?;
        let client_crypto =
                quinn::crypto::rustls::QuicClientConfig::try_from(Arc::new(client_crypto))?;
        endpoint.set_default_client_config(ClientConfig::new(Arc::new(client_crypto)));
        // twice, so as to allow 0-rtt to work on the second time
        for i in 0..2 {
            info!(%i, "client iteration");
            let connecting = endpoint.connect(
                "127.0.0.1:4433".to_socket_addrs().unwrap().next().unwrap(),
                "localhost",
            )?;
            // attempt to transmit 0-RTT data
            let duration = match connecting.into_0rtt() {
                Ok((conn, zero_rtt_accepted)) => {
                    debug!("attempting 0-rtt request");
                    let send_request_0rtt = send_request(&conn, num_bytes);
                    let mut send_request_0rtt_pinned = std::pin::pin!(send_request_0rtt);
                    tokio::select! {
                        result = &mut send_request_0rtt_pinned => result?,
                        accepted = zero_rtt_accepted => {
                            if accepted {
                                debug!("0-rtt accepted");
                                send_request_0rtt_pinned.await?
                            } else {
                                debug!("0-rtt rejected");
                                send_request(&conn, num_bytes).await?
                            }
                        }
                    }
                }
                Err(connecting) => {
                    debug!("not attempting 0-rtt request");
                    let conn = connecting.await?;
                    send_request(&conn, num_bytes).await?
                }
            };
            if i == 1 {
                println!("{}", duration.as_millis());
            }
            println!();
        }
        // tell the server to shut down so this process doesn't idle forever
        let _ = send_stop_server.send(());
        Ok(())
    }.instrument(info_span!("client"))));

    while tasks.join_next().await.is_some() {}
    Ok(())
}

async fn log_err<F: std::future::IntoFuture<Output=Result<(), Error>>>(task: F) {
    if let Err(e) = task.await {
        error!("{}", e);
    }
}